package edu.purdue.cs.ds.vss;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;

public class Reconstructor {
 
	public BigInteger reconstructSecret(
			ArrayList<ArrayList<BigInteger>> setOfShares) {

		ArrayList<BigInteger> shares = new ArrayList<BigInteger>();

		for (Iterator<ArrayList<BigInteger>> iterator = setOfShares.iterator(); iterator
				.hasNext();) {
			ArrayList<BigInteger> arrayList = (ArrayList<BigInteger>) iterator
					.next();

			for (Iterator<BigInteger> iterator2 = arrayList.iterator(); iterator2
					.hasNext();) {
				BigInteger tmpShare = (BigInteger) iterator2.next();
				if (!shares.contains(tmpShare)) {
					shares.add(tmpShare);
				}
			}
		}

		BigInteger result = new BigInteger("0");
		for (Iterator<BigInteger> iterator = shares.iterator(); iterator
				.hasNext();) {
			BigInteger tmp = (BigInteger) iterator.next();
			result = result.add(tmp).mod(Constants.q);
		}

		return result;
	}

	public BigInteger reconstructSecretFromStores(ArrayList<ShareStore> stores, 
			int n, int k) {
		
		//find the depth
		//Poke inside a secret and see the parent chain length
		ParentIndex pi = stores.get(0).getShares().get(0).getParentIndex();
		int depth = 1;
		while(pi.getParent() != null) {
			depth++;
			pi = pi.getParent();
		}
		ArrayList<ShareStore> tmpStores = stores;
		for(int i = 0; i < depth; i++) {
			stores = reconstruct(tmpStores, n);
			tmpStores = new ArrayList<ShareStore>();
			for(int j = 0; j < k+1; j++) { //Optimization we only need k+1
				tmpStores.add(stores.get(j));
			}
//			System.out.println("depth " + i + " done!");
			
		}
		
		return stores.get(0).getShares().get(0).getS();
	}

	private ArrayList<ShareStore> reconstruct(ArrayList<ShareStore> stores, int n) {

		ArrayList<ShareStore> tmpStores = new ArrayList<ShareStore>();
		for (int i = 0; i < n; i++) {
			tmpStores.add(new ShareStore());
		}

		ArrayList<HashMap<String, ArrayList<Share>>> serverIndexShareList = 
			new ArrayList<HashMap<String, ArrayList<Share>>>();
		
		for (int i = 0; i < n; i++) {
			
			HashMap<String, ArrayList<Share>> indexShareList = new HashMap<String, ArrayList<Share>>();
			serverIndexShareList.add(indexShareList);
		}
		
		for(int i = 0; i < stores.size(); i ++) {
			ShareStore store = stores.get(i);
			ArrayList<Share> shares = store.getShares();
			//Add these to the relevant list
			for (Iterator<Share> iterator = shares.iterator(); iterator.hasNext();) {
				Share share = iterator.next();
				ParentIndex pi = share.getParentIndex();
				int server = pi.getServer();
				//Get the server's list
				HashMap<String, ArrayList<Share>> indexShareList = serverIndexShareList.get(server);
				
				//Get the share list of the index
				ArrayList<Share> shareList = indexShareList.get(pi.toDisk());
				if(shareList == null) {
					shareList = new ArrayList<Share>();
					indexShareList.put(pi.toDisk(), shareList);
				}
				
				if(!shareList.contains(share)) {
					shareList.add(share);
				}
			}
		}

		//Now go through the each index share list and reconstruct
		for(int i = 0; i < n; i++) {
			HashMap<String, ArrayList<Share>> indexShareList = serverIndexShareList.get(i);
			Iterator<String> piIter = indexShareList.keySet().iterator();
			while (piIter.hasNext()) {
				String parentIndexStr = (String) piIter.next();
				ParentIndex parentIndex = new ParentIndex(parentIndexStr);
				//Try to reconstruct each index that the server holds
				ArrayList<Share> shareList = indexShareList.get(parentIndexStr);
				//How we have all shares of an index
				//reconstruct the secret with those
				if(!shareList.isEmpty()) {
					BigInteger result = new BigInteger("0");
					for (Iterator<Share> iterator = shareList.iterator(); iterator
							.hasNext();) {
						Share share =  iterator.next();
						result = result.add(share.getS()).mod(Constants.q);
					}
					//Create the share and add it to store
					Share tmpShare = new Share(result, parentIndex.getIndex(), parentIndex.getParent());
					tmpStores.get(i).update(tmpShare);
				}
			}
			
		}
		
		return tmpStores;
	}

}
